# -*- coding:utf-8 -*-
import os
import time
import sys
import subprocess


def upgrade(loadertool_path, local_image_path):
	upgrade_loader_cmd = "%s UL %s\MiniLoaderAll.bin -noreset" % (loadertool_path, local_image_path)
	h = exec(upgrade_loader_cmd)
	if "Upgrade loader ok" not in h:
		print("Download MiniLoaderAll.bin Fail!")
		return False
	else:
		print("Download MiniLoaderAll.bin Success!")
		time.sleep(3)
		write_gpt_cmd = "%s DI -p %s\parameter.txt" % (loadertool_path, local_image_path)
		j = exec(write_gpt_cmd)
		if "Write gpt ok" not in j:
			print("Failed to execute the parameter.txt")
			return False
		else:
			print("Successfully executed parameter.txt.")
			time.sleep(5)
			download_uboot_cmd = "%s DI -uboot %s\\uboot.img %s\parameter.txt" % (loadertool_path, local_image_path, local_image_path)
			k = exec(download_uboot_cmd)
			if "Download image ok" not in k:
				print("Failed to download the uboot.image!")
				return False
			else:
				print("The uboot.image downloaded successfully!")
				time.sleep(5)
				if not flash_version(loadertool_path, local_image_path):
					return False
				return True


def flash_version(loadertool_path, local_image_path):
	partList = ["boot_linux", "system", "vendor", "userdata", "resource", "ramdisk", "chipset", "sys-prod", "chip-prod", "updater"]
	for i in partList:
		if not os.path.exists("%s/%s.img" % (local_image_path, i)):
			print("%s.img is not exist, ignore" % i)
			continue
		loadcmd = "%s DI -%s %s/%s.img" % (loadertool_path, i, local_image_path, i)
		p = exec(loadcmd)
		if "Download image ok" not in p:
			print("try download %s again!" % i)
			time.sleep(3)
			second_cmd = "%s DI -%s %s/%s.img" % (loader_tool_path, i, local_image_path, i)
			f = exec(second_cmd)
			print(f)
			if "Download image ok" not in f:
				print("Failed to download the %s.img!" % i)
				return False
			return True
		else:
			print("The %s.img downloaded successfully!" % i)
	return True


hdc = 'hdc_std'


def connect():
	targets = subprocess.check_output(['hdc_std', 'list', 'targets']).decode('utf-8').splitlines()
	targets = list(filter(lambda s: s != '[Empty]', targets))
	if len(targets) > 1:
		for i in range(len(targets)):
			print('%d: %s' % (i+1, targets[i]))
		while True:
			try:
				choice = targets[int(input("choose which device to burn (1-%d): " % len(targets))) - 1]
			except BaseException:
				print("bad choice! please choose again.\n")
			else:
				global hdc
				hdc = '%s -t %s' % (hdc, choice)
				print("hdc command is", hdc)
				break
	for i in range(5):
		ret = exec(hdc + " shell exit")
		print(ret)
		if "Fail" not in ret:
			return True
		os.system(hdc + " kill")
		time.sleep(1)
		os.system(hdc + " start")
		time.sleep(1)
	return False


def exec(cmd):
	result = "".join(os.popen(cmd).readlines())
	return result


if __name__ == "__main__":
	# 修改配置
	# upgrade_tool.exe的路径
	loadertool_path = r"upgrade_tool.exe"
	# 解压后的RK镜像目录
	#local_image_path = r'D:\MobileUpgrade\local_img\aa\img'
	if len(sys.argv) <= 1:
		exit(input("images path must be specified!\n"))
	if sys.argv[1] == '?':
		exit(print('py文件关联测试成功!\n'))
	local_image_path = sys.argv[1]
	print(local_image_path, os.getcwd())

	# 检查hdc是否连接，并尝试重连5次
	print("--------------检查hdc是否连接--------------")
	connect_status = connect()
	if not connect_status:
		print("---------连接失败，尝试直接刷机---------")
	else:
		print("---------连接成功---------")
	# 进入loader模式
	print("---------进入loader模式----------")
	os.system(hdc + " shell reboot loader")
	time.sleep(5)
	# 刷机
	print("---------刷机开始---------")
	upgrade_status = upgrade(loadertool_path, local_image_path)
	if upgrade_status:
		print("---------刷机成功---------")
	else:
		print("---------刷机失败---------")
		exit(1)
	time.sleep(5)
	# 重启
	os.system(f"{loadertool_path} RD")
	#time.sleep(30)
	input("Done! Press Enter to exit..")

